import numpy as np
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
import parl


# 灰度处理后的Model
class HKPool2DModel(parl.Model):
    def __init__(self, act_dim, obs_dim, dueling=False):
        super(HKPool2DModel, self).__init__()
        self.dueling = dueling
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        # 第一个卷积层, 高为obs_dim[0], 宽为obs_dim[1], 4, 卷积核大小为11, 步长为4, padding为0
        kernel_size = 11
        stride = 4
        padding = 3
        self.conv1 = nn.Conv2D(4, 16, kernel_size, stride, padding)
        # 计算卷积后的输出维度
        h_out_dim, w_out_dim = self.out_dim(obs_dim[0], obs_dim[0], kernel_size, stride, padding)

        # 第二个卷积层, 高为h_out_dim, 宽为w_out_dim, 通道数为32, 卷积核大小为5, 步长为1, 填充为2
        kernel_size = 5
        stride = 1
        padding = 2
        self.conv2 = nn.Conv2D(16, 32, kernel_size=kernel_size, stride=stride, padding=padding)
        # 计算卷积后的输出维度
        h_out_dim, w_out_dim = self.out_dim(h_out_dim, w_out_dim, kernel_size, stride, padding)

        # 第三个卷积层, 高为h_out_dim, 宽为w_out_dim, 通道数为32, 卷积核大小为3, 步长为1, 填充为1
        kernel_size = 3
        stride = 1
        padding = 1
        self.conv3 = nn.Conv2D(32, 32, kernel_size=kernel_size, stride=stride, padding=padding)
        # 计算卷积后的输出维度
        h_out_dim, w_out_dim = self.out_dim(h_out_dim, w_out_dim, kernel_size, stride, padding)

        # 池化层
        self.pool = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
        # 计算池化后的输出维度
        h_out_dim, w_out_dim = self.out_dim(h_out_dim, w_out_dim, 2, 2, 0)

        # 三个全连接层, 将卷积后的输出转换为act_dim维度的输出, 作为Q值
        feature_dim = 1024
        self.fc1 = nn.Linear(h_out_dim * w_out_dim * 32, feature_dim)
        # self.fc1 = nn.Linear(1811328, 128)
        self.fc2 = nn.Linear(feature_dim, feature_dim // 2)
        self.fc3 = nn.Linear(feature_dim // 2, act_dim)
        if self.dueling:
            self.fc4 = nn.Linear(feature_dim // 2, 1)
        self.flatten = nn.Flatten()

    # 计算卷积后的输出维度
    def out_dim(self, h_in_dim, w_in_dim, kernel_size, stride, padding):
        h_out_dim = (h_in_dim - kernel_size + 2 * padding) // stride + 1
        w_out_dim = (w_in_dim - kernel_size + 2 * padding) // stride + 1
        return h_out_dim, w_out_dim

    def forward(self, obs):
        h1 = F.relu(self.conv1(obs))  # 卷积
        h2 = F.relu(self.conv2(h1))  # 卷积
        h3 = F.relu(self.conv3(h2))  # 卷积
        h_pool = self.pool(h3)  # 池化
        h4 = self.flatten(h_pool)  # 将卷积后的输出展平
        if self.dueling:
            h5 = F.relu(self.fc1(h4))
            h6 = F.relu(self.fc2(h5))
            adv = self.fc3(h6)
            value = self.fc4(h6)
            Q = value + adv - paddle.mean(adv, axis=1, keepdim=True)
        else:
            h5 = F.relu(self.fc1(h4))  # 全连接
            h6 = F.relu(self.fc2(h5))  # 全连接
            Q = self.fc3(h6)  # 全连接
        return Q


if __name__ == '__main__':
    model = HKPool2DModel(4, (720, 720))
    obs = np.zeros((32, 4, 720, 720), dtype=np.float32)
    obs = paddle.to_tensor(obs)
    Q = model(obs)
    print(Q)
